Add const shadows for forward mode AD on RegionBranchOpInterface#2780
Merged
Conversation
vimarsh6739
force-pushed
the
scf.for-fix2
branch
2 times, most recently
from
51333d0 to
450fff0
Compare
wsmoses
reviewed
May 5, 2026
vimarsh6739
force-pushed
the
scf.for-fix2
branch
2 times, most recently
from
2fcb520 to
d31bb21
Compare
wsmoses
reviewed
May 7, 2026
wsmoses
reviewed
May 7, 2026
| // forceAugmentedReturns will not shadow const operands. No need to add | ||
| // to the invertPointers map since `operand` is const, the shadow will | ||
| // be unused. | ||
| for (const RegionSuccessor &successor : entrySuccessors) { |
Member
There was a problem hiding this comment.
I don't think this should be here, but within the impl of createWithShadows, right?
Member
Author
There was a problem hiding this comment.
No, it has to be here. We are shadowing only the const operands for any augmented newOp.
Member
Author
There was a problem hiding this comment.
createWithShadows correctly adds const shadows. But newOp doesn't have them(since it was created with forceAugmentedReturns), so the takeBody() discards the args in the replacement op
We have 2 choices here - either add the const shadows to newOp or the replacement (after takeBody). The end result is the same.
vimarsh6739
force-pushed
the
scf.for-fix2
branch
6 times, most recently
from
736e317 to
a0a2d7b
Compare
vimarsh6739
changed the title
RegionBranchOpInterface
When an operand has const activity, we may still need to create a shadow for it (particularly in scf.for, where a const iter_arg may still have to be shadowed if the result and the terminator are active)
wsmoses
approved these changes
Jun 5, 2026
wsmoses
added a commit
that referenced
this pull request
Jun 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This primarily affects scf.for and affine.for, but should also handle a the same case in scf.parallel and affine.parallel ops (iter-arg is a constant/dead inside loop but terminators still have activity).
Aside: I think one interesting thing we can do here is to reuse the successor information to effectively eliminate the creation of some dead regions. This already kind of happens in scf.for which prunes the entry successors depending on if the loop has no iterations, or non-zero trip count(e.g. if the only entry successor is the parent, then just skip diffing the regionBranchOp body).